import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
print(sys.path)
# sys.exit(0)
from argparse import ArgumentParser
from ultralytics import YOLO
import os
# os.environ['CUDA']
# print(os.getenv('CUDA_VISIBLE_DEVICES'))
parser = ArgumentParser()
if os.getenv('face_dev'):
    parser.add_argument('--weights', default='/root/code/cvmark/face/project/train/models/face_det/train/weights/last.pt')
else:
    parser.add_argument('--weights', default='/project/train/models/face_det/train/weights/best.pt')
parser.add_argument('--imgsz', default=960, type=int)
args = parser.parse_args()
weights =args.weights
imgsz = args.imgsz
args = parser.parse_args()
yolo = YOLO(args.weights)
yolo.export(format='engine', nms=True, imgsz=imgsz, int8=False, half=False, device=0)
# 
# yolo.export(format='onnx', nms=True, imgsz=1920, int8=False)
# replace .pt with .onnx, with the other path the same
# onnx_weights = os.path.basename(args.weights).replace('.pt', '.onnx')
# print(onnx_weights)
# yolo_onnx = YOLO()
